from flask import Flask, request, jsonify, send_from_directory
from flask_cors import CORS, cross_origin
import os
from sparkai.llm.llm import ChatSparkLLM, ChunkPrintHandler
from sparkai.core.messages import ChatMessage

app = Flask(__name__, static_folder='./')
CORS(app)  # 允许跨域请求

# 星火认知大模型配置信息
SPARKAI_URL = 'wss://spark-api.xf-yun.com/v3.5/chat'
SPARKAI_APP_ID = 'd16f688b'
SPARKAI_API_SECRET = 'M2Q2ZDM3MGY4YzZmYzgyM2NiNDdmNjAx'
SPARKAI_API_KEY = 'f5871f26e082699875dcdc77ce2cb0f9'
SPARKAI_DOMAIN = 'generalv3.5'

# 初始化星火大模型
spark = ChatSparkLLM(
    spark_api_url=SPARKAI_URL,
    spark_app_id=SPARKAI_APP_ID,
    spark_api_key=SPARKAI_API_KEY,
    spark_api_secret=SPARKAI_API_SECRET,
    spark_llm_domain=SPARKAI_DOMAIN,
    streaming=False
)

@app.route('/api/chat', methods=['POST'])
@cross_origin()
def chat():
    data = request.get_json()
    message = data.get('message', 'Hello')

    # 调用星火大模型生成回复
    messages = [ChatMessage(role="user", content=message)]
    handler = ChunkPrintHandler()
    
    # 调用大模型生成结果
    response = spark.generate([messages], callbacks=[handler])
    
    # 提取生成的文本内容
    generated_text = response.generations[0][0].text if response.generations else "No response from model"

    # 获取请求来源的 IP 地址
    remote_addr = request.headers.get('X-Forwarded-For', request.remote_addr)

    # 将用户消息和 AI 回复分别写入文件，AI 回复的 IP 为 "0.0.0.0"
    with open('chat_log.txt', 'a') as file:
        # 记录用户消息
        file.write(f"IP: {remote_addr}, Message: {message}\n")
        # 记录AI回复
        file.write(f"IP: 0.0.0.0, Message: {generated_text}\n")

    # 返回大模型生成的文本给前端
    return jsonify({'greeting': generated_text})

@app.route('/api/chat_log', methods=['GET'])
@cross_origin()
def chat_log():
    try:
        # 读取聊天记录文件的最后50行
        with open('chat_log.txt', 'r') as file:
            lines = file.readlines()
            # 如果聊天记录不足50行，则返回所有记录
            last_50_lines = lines[-50:] if len(lines) >= 50 else lines
        # 返回聊天记录的列表
        return jsonify(last_50_lines)
    except FileNotFoundError:
        # 如果文件不存在，返回空列表
        return jsonify([])

@app.route('/', methods=['GET'])
def gethtml():
    return send_from_directory(app.static_folder, './chat.html')

if __name__ == '__main__':
    app.run(debug=True, host='0.0.0.0', threaded=True)  # 运行应用，并开启调试模式